import genesis as gs
import numpy as np
import torch

gs.init(backend=gs.gpu)

viewer_options = gs.options.ViewerOptions(
    camera_pos=(0, -3.5, 2.5),
    camera_lookat=(0.0, 0.0, 0.5),
    camera_fov=40,
    max_FPS=30,
)

scene = gs.Scene(
    viewer_options=viewer_options,
    sim_options=gs.options.SimOptions(
        dt=0.01,
        substeps=4, # 增加子步骤提高抓取稳定性
        gravity=(0, 0, -10),
    ),
    show_viewer=True,
    renderer=gs.renderers.Rasterizer(),
)

ground = scene.add_entity(
    gs.morphs.Plane()
)

franka = scene.add_entity(
    gs.morphs.MJCF(
        file="xml/franka_emika_panda/panda.xml",
        pos = (0, 0, 0),
        euler = np.random.rand(3) * 2 * np.pi,
        scale = 1.0,          
    ),
)

sphere_radius = 0.03
sphere_pos = np.array([0.65, 0.0, sphere_radius])

target_sphere = scene.add_entity(
    gs.morphs.Sphere(
        radius=sphere_radius,
        pos=sphere_pos
    ),
)

scene.build()

joint_names = [
    "joint1",
    "joint2",
    "joint3",
    "joint4",
    "joint5",
    "joint6",
    "joint7",
    "finger_joint1",
    "finger_joint2",
]

dofs_idx = [franka.get_joint(name).dof_idx_local for name in joint_names]



############ 可选：设置控制增益 ############
# 设置位置增益
franka.set_dofs_kp(
    kp             = np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
    dofs_idx_local = dofs_idx,
)
# 设置速度增益
franka.set_dofs_kv(
    kv             = np.array([450, 450, 350, 350, 200, 200, 200, 10, 10]),
    dofs_idx_local = dofs_idx,
)
# 设置安全的力范围
franka.set_dofs_force_range(
    lower          = np.array([-87, -87, -87, -87, -12, -12, -12, -100, -100]),
    upper          = np.array([ 87,  87,  87,  87,  12,  12,  12,  100,  100]),
    dofs_idx_local = dofs_idx,
)

hand = franka.get_link("hand")

loops = 20

# 随机生成目标位置
target_pos = np.random.rand(3)
target_quat = np.random.rand(4)

src_pos = np.array([0.5,0,0.2])
src_quat = np.array([0,1,0,0])

for i in range(loops):
    # 计算抓取方向
    qpos = franka.inverse_kinematics(
        link = hand,
        pos = target_pos,
        quat = target_quat,    
    )
    path = franka.plan_path(
        qpos_goal = qpos,
        num_waypoints = 200,
    )
    if len(path) == 0:
        continue
    for waypoint in path:
        franka.control_dofs_position(waypoint, dofs_idx)
        scene.step()
    for z in range(100):
        scene.step()
    
    # 恢复
    # 计算抓取方向
    qpos = franka.inverse_kinematics(
        link = hand,
        pos = src_pos,
        quat = src_quat,    
    )
    path = franka.plan_path(
        qpos_goal = qpos,
        num_waypoints = 200,
    )
    for waypoint in path:
        franka.control_dofs_position(waypoint, dofs_idx)
        scene.step()
    for z in range(100):
        scene.step()